# -*- coding: utf-8 -*-
import numpy as np
import matplotlib.pyplot as plt 
import arcpy
import math
from matplotlib.ticker import MultipleLocator
from matplotlib.ticker import FormatStrFormatter
from matplotlib.font_manager import FontProperties
font_set1 = FontProperties(fname=r'C:\windows\fonts\simsun.ttc',size=12)

def laa(x):
    x_mean = np.mean(x)
    lxx = np.sum((x-x_mean)**2)
    return lxx

#用于计算Lxy
def lab(x,y):
    x_mean = np.mean(x)
    y_mean = np.mean(y)
    lxy = np.sum((x-x_mean)*(y-y_mean))
    return lxy

def polyfit_one(x, y, alpha):
    if len(set(x))==1:
        if len(set(y))==1:
            poly_val=[999999,999999]
            R=999999
            R2=999999
            pval=999999
            linear_test=999999
            poly_int=999999
            lxx=999999
            lyy=999999
            lxy=999999
            sigma_est=999999
            x_mean=999999
            y_mean=999999
            test_level=999999
            n=999999
        else:
            poly_val=[99999999,np.mean(x)]
            R=999999
            R2=999999
            pval=999999
            linear_test=999999
            poly_int=999999
            lxx=999999
            lyy=999999
            lxy=999999
            sigma_est=999999
            x_mean=999999
            y_mean=999999
            test_level=999999
            n=999999
    else:
        assert len(x) == len(y)
        n = len(x)
        assert n > 2
        lxx = laa(x)
        lyy = laa(y)
        lxy = lab(x, y)
        R = lxy/(np.sqrt(lxx) * np.sqrt(lyy))
        R2 = R*R   #计算相关系数与决定系数
        print(R2)
        b_est = lxy/lxx  #计算b估计
        x_mean = np.mean(x)
        y_mean = np.mean(y)
        a_est = y_mean - b_est * x_mean   #计算a估计
        poly_val = (a_est, b_est)

    #返回回归模型相应参数
    test_val = {'R': R,
                'R2': R2
                }
    process_val = {'lxx': lxx,
                   'lyy': lyy,
                   'lxy': lxy,
                   'x_mean': x_mean,
                   'y_mean': y_mean,
                   'ndim': n
                   }
    return (poly_val, test_val, process_val)




x=[]
y=[]



fc = arcpy.GetParameterAsText(0)

field1 = "TB"
field2 = "E"
cursor = arcpy.SearchCursor(fc)
for row in cursor:
    x.append(row.getValue(field1))
    y.append(row.getValue(field2))


fig, ax = plt.subplots(figsize = (15,2+int((max(np.array(y))+200-min(np.array(y)))/5.33)))
fig.patch.set_facecolor('white')
ax.xaxis.set_ticks_position('bottom')
ax.spines['bottom'].set_position(('data', 0))
ax.yaxis.set_ticks_position('left')
ax.spines['left'].set_position(('data', 0))
ax.spines['right'].set_color('none')
ax.spines['top'].set_color('none')
    
ax.grid(ls='--',linewidth=0.4)
#设置间距
xmajorLocator = MultipleLocator(1)
xmajorFormatter = FormatStrFormatter('%d')
ax.xaxis.set_major_locator(xmajorLocator)
ax.xaxis.set_major_formatter(xmajorFormatter)

if max(np.array(y))>700:
    ymajorLocator = MultipleLocator(5.33)
    y_maxlim=max(np.array(y)) +100
    y_minlim=min(np.array(y)) -100
else:
    ymajorLocator = MultipleLocator(5.33)
    y_maxlim=max(np.array(y)) +100
    y_minlim=min(np.array(y)) -100
ymajorFormatter = FormatStrFormatter('%1.1f')
ax.yaxis.set_major_locator(ymajorLocator)
ax.yaxis.set_major_formatter(ymajorFormatter)


if 0<len(x)<=2:
    ax.scatter(x, y, s=15,color='r')
else:
    poly_val,test_val,process_val = polyfit_one(np.array(x),np.array(y),0.05)
    if poly_val[0]==999999:
        ax.scatter(x,y,s=10,color='r')
    elif poly_val[0]==99999999:
        poly_text1 = "$E = %s*TB"%(str(round(poly_val[1],2)))
#         poly_text1='E='+str(round(poly_val[1],2))+'*TB'
        X_test = np.ones(1000000)*poly_val[1]
        Y_test = np.linspace(-math.ceil(max(y))-50,math.ceil(max(y))+50,1000000)
        plt.plot(X_test,Y_test,'-',color='darkred',label=poly_text1,linewidth=1)
        ax.scatter(x,y,s=10,color='r')
    elif poly_val[0]!=99999999 and poly_val[1]!=0:
        poly_text1 = "$E = %s*TB  + %s$($R^2 = %s$)"%(str(round(poly_val[1],2)),str(round(poly_val[0],2)),str(round(test_val['R2'],2)))#加$会使字体倾斜加粗
#         poly_text1='E='+str(round(poly_val[1],2))+'*TB+'+str(round(poly_val[0],2))+'(R^2='+str(round(test_val['R2'],2))+')'
        ax.scatter(x,y,s=10,color='r')
        X_test = np.linspace(0,math.ceil(max(x))+6, 1000000)
        Y_test = X_test*poly_val[1]+poly_val[0]
        plt.plot(X_test,Y_test,'-',color='darkred',label=poly_text1,linewidth=1)
    elif poly_val[1] == 0:
#         poly_text1=str(round(poly_val[0],2))+'(R^2='+str(round(test_val['R2'],2))+')'
        poly_text1 = "$E = %s$" %(str(round(poly_val[0],2)))
        ax.scatter(x,y,s=10,color='r')
        X_test = np.linspace(0,math.ceil(max(x))+6, 1000000)
        Y_test = X_test* poly_val[1]+poly_val[0]
        plt.plot(X_test,Y_test,'-',color='darkred',label=poly_text1,linewidth=1)
           
ax.plot(np.linspace(0, 30, 1000000), 2*5.893*1.8099*np.linspace(0, 30, 1000000), '--k',linewidth=1,alpha = 0.5)
ax.plot(np.linspace(0, 30, 1000000), 5.893*1.8099*np.linspace(0, 30, 1000000), '--k', linewidth=1,alpha = 0.5)
ax.plot(np.linspace(0, 30, 1000000), (1.0/3)*5.893*1.8099*np.linspace(0, 30, 1000000), '--k',linewidth=1,alpha = 0.5)
ax.plot(np.linspace(0, 30, 1000000), (-1.0/3)*5.893*1.8099*np.linspace(0, 30, 1000000), '--k',linewidth=1,alpha = 0.5)
ax.plot(np.linspace(0, 30, 1000000), -0.5*5.893*1.8099*np.linspace(0, 30, 1000000), '--k',linewidth=1,alpha = 0.5)
ax.plot(np.linspace(0, 30, 1000000), -5.893*1.8099*np.linspace(0, 30, 1000000), '--k',linewidth=1,alpha = 0.5)
        
        

plt.xlim(0,30)
plt.ylim(y_minlim,y_maxlim)
legend = ax.legend(loc="upper center",prop=font_set1)

plt.show()
import time
time.sleep(2)
plt.close()